import os, sys
import numpy as np
from scipy.stats import sigmaclip
from astropy.io import fits
from astropy.table import Table
from skimage.measure import label


# mask bit definition
# warm pixels (5 sigma) = 2
# hot pixels (bad) = 32
# dark pixel (< 50%) = 64
# charge trap = 128
# no data = 512

def merge_mask(flist, fout, band=None):
    n = len(flist)
    mask = None
    for i in range(n):
        if i == 0:
            mask = fits.getdata(flist[i]).astype('>i2')
            #print("masktype",mask.dtype,flist[i])
        else:
            #print("masktype",fits.getdata(flist[i]).dtype,flist[i])
            mask = np.bitwise_or(mask, fits.getdata(flist[i]).astype('>i2'))
    hdu = fits.PrimaryHDU(mask.astype(np.int16))
    if band is not None:
        hdu.header['BAND'] = band
    for i in range(n):
        hdu.header['FMASK{}'.format(i+1)] = os.path.basename(flist[i])
    hdu.writeto(fout, overwrite=True)


def get_dmask(fdark, fout):
    with fits.open(fdark) as img:
        dark = img[0].data

        # invalid pixels = 512
        valid = np.isfinite(dark)
        output = (~valid).astype(int) * 512

        # hot pixels (bad) = 32
        index = dark > 25
#        index = dark > 0.08
        output[index] = 32

        # warm pixels = 2
        if 'WARM_PIXEL_DC' in os.environ:
            warm_pixel_dc = float(os.environ['WARM_PIXEL_DC'])
            print('==> get warm_pixel_dc = {} from environment variable'.format(warm_pixel_dc))
            # sys.exit(0)
            if warm_pixel_dc > 0:
                # clip = sigmaclip(dark[valid], 4, 4)[0]
                w_index = dark > warm_pixel_dc
                output[w_index] = output[w_index] + 2   # hot pixels are also warm pixels !!

        # cosmic ray = 1024
        
        # output
        hdu = fits.PrimaryHDU(output.astype(np.int16), header=img[0].header)
        hdu.writeto(fout, overwrite=True)
        print('dark mask written: {}'.format(fout))


def _moving_sum(data, box):
    ny, nx = data.shape
    dd = np.zeros((ny + box * 2, nx + box * 2), dtype=data.dtype)
    dd[box:-box, box:-box] = data
    dd = np.cumsum(dd, axis=0)
    dd = dd[box:, :] - dd[:-box, :]
    dd = np.cumsum(dd, axis=1)
    dd = dd[:, box:] - dd[:, :-box]
    return dd[box // 2:-box // 2, box // 2:-box // 2]


def get_smooth_flat(flat, boxsize=100):
    dd = flat.copy()
    med = np.nanmedian(dd)
    good = np.logical_and(dd < med * 1.2, dd > med * 0.8)  # flag outlier
    dd[~good] = 0
    dd = _moving_sum(dd, boxsize)
    nn = _moving_sum(good.astype(int), boxsize)
    index = nn == 0
    nn[index] = 10000
    flat_smth = dd / nn
    return flat_smth


def get_fmask(fflat, fout):
    with fits.open(fflat) as img:
        flat = img[0].data

        # invalid pixels = 512
        valid = np.isfinite(flat)
        output = (~valid).astype(np.int16) * 512

        # dark pixel (< 50%) = 64
        flat_smth = get_smooth_flat(flat)
        index = flat < (flat_smth * 0.5)
        output[index] = 64

        # output
        hdu = fits.PrimaryHDU(output.astype(np.int16), header=img[0].header)
        hdu.header['FLATFILE'] = os.path.basename(fflat)
        hdu.writeto(fout, overwrite=True)
        print('flat mask written: {}'.format(fout))


def mark_badpixel(msk):
    ny, nx = msk.shape
    bad = (msk >= 32) * (msk < 512)  # hot pixel (32), dark pixel (64) and charge trap (128)
    output = np.zeros((ny, nx), dtype=np.uint8)
#
    # bad column = 2
    for i in range(nx - 4):
        if bad[2:-2, i + 2].sum() > 100:
            output[2:-2, i + 2] = output[2:-2, i + 2] + 2
    # bad row = 4
    for i in range(ny - 4):
        if bad[i + 2, 2:-2].sum() > 100:
            output[i + 2, 2:-2] = output[i + 2, 2:-2] + 4

    # individual bad pixel = 8
    index = np.logical_and(output == 0, bad)
    output[index] = 8

    # charge traps = 16
    index = np.logical_and((msk & 32) == 0, (msk & 64) == 0)
    index = np.logical_and((msk & 128) > 0, index)  # remove hot pixel and dark pixel
    output[index] = output[index] + 16

    # image edge = 1
    output[:, 0:2] = 1
    output[:, -2:] = 1
    output[0:2, :] = 1
    output[-2:, :] = 1
    # invalid data = 1
    index = (msk & 512) > 0
    output[index] = 1

#    #revised @ 20230413 to take warm column as defect
#    bad = (msk>=2) * (msk<=512)
#    #bad column = 2
#    for i in range(nx - 4):
#        if bad[2:-2, i + 2].sum()>100:
#            output[2:-2, i + 2] = output[2:-2, i + 2] + 2
#    #bad row = 4
#    for i in range(ny - 4):
#        if bad[i + 2, 2:-2].sum() > 100:
#            output[i + 2, 2:-2] = output[i + 2, 2:-2] + 4

    return output


def get_largest_cluster(mask):
    segs = label(mask)
    seg, npix = np.unique(segs, return_counts=True)
    #print("seg",seg,"npix",npix)
    index = seg > 0
    tab = Table()
    xcs = []
    ycs = []
    xss = []
    yss = []
    #20230318 add cluster table generation
    if index.sum() > 0:
        npixs = npix[index]        
        seg_idx = seg[index]
        seg_idx = seg_idx[npixs.argsort()]
        npixs.sort()
        #print("npixs",npixs)
        for i in range(len(npixs)):
            #print(seg_idx[-(i+1)],npixs[-(i+1)])
            c = np.where(segs==seg_idx[-(i+1)])
            #print(c)
            xmax = np.max(c[1])
            xmin = np.min(c[1])
            ymax = np.max(c[0])
            ymin = np.min(c[0])
            xc = int((xmin+xmax)/2)
            yc = int((ymin+ymax)/2)
            xs = int((xmax-xmin))
            ys = int((ymax-ymin))
            xcs.append(xc)
            ycs.append(yc)
            xss.append(xs)
            yss.append(ys)
        tab['npix'] = npixs[::-1]
        tab['x_center'] = xcs
        tab['y_center'] = ycs
        tab['x_size'] = xss
        tab['y_size'] = yss
        result = npixs[-1]
        #print("largest_cluster",int((xmin1+xmax1)/2),int(xmax1-xmin1),int((ymin1+ymax1)/2),int(ymax1-ymin1))
        #print("second_cluster",int((xmin2+xmax2)/2),int(xmax2-xmin2),int((ymin2+ymax2)/2),int(ymax2-ymin2))
    else:
        print("no cluster!")
        result = 1
    return result,tab


def count_defects(bad):
    
    valid = bad != 1
    pix_tot = valid.sum()
    # column
    index = np.logical_and((bad & 2) > 0, valid)
    col = index.sum() / pix_tot
    # row
    index = np.logical_and((bad & 4) > 0, valid)
    row = index.sum() / pix_tot
    # individual
    index = np.logical_and((bad & 8) > 0, valid)
    indiv = index.sum() / pix_tot
    #2023/03/18 add return tab_cluster
    npix_cluster,tab_cluster = get_largest_cluster(index)  # clustered individual

    # charge traps
    # index = np.logical_and((bad & 16) > 0, valid)
    index = np.logical_and(bad == 24, valid)
    trap = index.sum() / pix_tot

    tab = Table()
    tab['npix'] = [pix_tot, ]
    tab['col'] = np.float32(col)
    tab['row'] = np.float32(row)
    tab['lines'] = tab['col'] + tab['row']
    tab['indiv'] = np.float32(indiv)
    tab['trap'] = np.float32(trap)
    tab['npix_cluster'] = npix_cluster

    return tab, tab_cluster
